#include <windows.h>
#include <iostream>

#define THREAD 12


// 多线程矩阵乘法,windowsAPI实现
typedef struct {
    int start;
    int end;
    int k;
    float *a;
    float *b;
    float *res;
}Pot;


DWORD WINAPI matmul_t(LPVOID lpParam){
    Pot *p = (Pot *)lpParam;
    
    int start = p->start;
    int end = p->end;
    int k = p->k;
    float *A = p->a;
    float *B = p->b;
    float *C = p->res;

    for(int i = start; i < end; i++) {
        int row = i / k;
        int col = i % k;
        for (int j = 0; j < k; j++) {
            C[i] += A[row * k + j] * B[j * k + col];
        }
    }
    
     return 0;
}


int matmul(float *A, float *B, float *C, int m, int n, int k) {
    
    int nums = m * n; // 矩阵元素个数
    int thread_num = THREAD; // 线程数量
    
    HANDLE *hThreads = new HANDLE[thread_num];
    Pot *p = new Pot[thread_num];

    // 线程分配
    int blocksdim;
    if (nums%thread_num != 0) {
        if (nums < thread_num) {
            for (int i = 0; i < thread_num; i++) {
                int end;
                int start;
                if (i < nums) {
                    end = i + 1;
                    start = i;
                }else{
                    end = nums;
                    start = nums;
                }
                
                p[i] = Pot{start, end, k, A, B, C};
            }
        }else{
            blocksdim = nums/(thread_num -1); // 计算每个线程需要处理的数据量
            for (int i = 0; i < thread_num; i++) {
                int end;
                if (i == thread_num - 1) {
                    end = nums;
                }
                else{
                    end = blocksdim * (i + 1);
                }
                int start = blocksdim * i;
                
                p[i] = Pot{start, end, k, A, B, C};
            }
        }
    }else{
        blocksdim = nums/thread_num;

        for (int i = 0; i < thread_num; i++) {
                int end = blocksdim * (i + 1);
                int start = blocksdim * i;
                p[i] = Pot{start, end, k, A, B, C};
            }
    }

    
    // 创建线程
    for (int i = 0; i < thread_num; i++) {
            hThreads[i] = CreateThread(NULL, 0, matmul_t, &p[i], 0, NULL);
            if (hThreads[i] == NULL) {
                std::cerr << "CreateThread failed (" << GetLastError() << ")" << std::endl;
            delete[] hThreads;
            delete[] p;
            return 1;
            }
    }

    WaitForMultipleObjects(thread_num, hThreads, TRUE, INFINITE);

    for (int i = 0; i < thread_num; i++) {
        CloseHandle(hThreads[i]);
        
    }

    delete[] hThreads;
    delete[] p;

    return 0;
}

int main() {
    int m = 10000;
    int n = 10000;
    int k = 10000;
    float A[m * k] = {1.0f};
    float B[k * n] = {1.0f};
    float C[m * n] = {0.0f};

    matmul(A, B, C, m, n, k);

    // for(int i = 0; i < m; i++) {
    //     for(int j = 0; j < n; j++) {
    //         std::cout << C[i * n + j] << " ";
    //     }
    //     std::cout << std::endl;
    // }
    return 0;

}